package org.zjx.server;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.message.messages.RouteResponse;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

public class RouteManager {
    private static final Logger logger = LoggerFactory.getLogger(RouteManager.class);
    private final Map<String, WorkerInfo> workers = new ConcurrentHashMap<>();
    private final Map<String, String> serviceToWorker = new ConcurrentHashMap<>();

    public void registerWorker(WorkerInfo worker) {
        workers.put(worker.getId(), worker);
        worker.getServices().forEach(service ->
                serviceToWorker.put(service, worker.getId())
        );
        logger.info("Worker registered: {} with services: {}", worker.getId(), worker.getServices());
    }

    public void removeWorker(String workerId) {
        WorkerInfo worker = workers.remove(workerId);
        if (worker != null) {
            worker.getServices().forEach(serviceToWorker::remove);
        }
        logger.warn("Worker removed: {}", workerId);
    }

    public RouteResponse queryRoute(String serviceName) {
        String workerId = serviceToWorker.get(serviceName);
        if (workerId == null) {
            logger.warn("Service not found: {}", serviceName);
            return RouteResponse.unavailable(
                    serviceName,
                    "Service not registered in cluster"
            );
        }

        WorkerInfo worker = workers.get(workerId);
        if (worker == null) {
            logger.error("Worker not available: {}", workerId);
            return RouteResponse.unavailable(
                    serviceName,
                    "Worker not available"
            );
        }

        logger.info("Route found for {} -> {}", serviceName, workerId);
        // 重定向到目标节点
        return RouteResponse.redirect(
                serviceName,
                worker.getHost(),
                worker.getPort(),
                workerId
        );
    }

    // 更新worker心跳时间
    public void updateHeartbeat(String workerId) {

        if (workerId == null || workerId.isEmpty()) {
            logger.warn("Invalid workerId in heartbeat");
            return;
        }

        WorkerInfo worker = workers.get(workerId);
        if (worker != null) {
            worker.setLastHeartbeat(System.currentTimeMillis());
            logger.debug("Heartbeat updated for worker: {}", workerId);
        } else {
            logger.warn("Received heartbeat from unknown worker: {}", workerId);
        }
    }

    // 获取所有worker信息（用于心跳检测）
    public Map<String, WorkerInfo> getWorkerInfos() {
        return Collections.unmodifiableMap(workers);
    }

    public static class WorkerInfo {
        private final String id;
        private final String host;
        private final int port;
        private volatile long lastHeartbeat;
        private final Set<String> services;

        public WorkerInfo(String id, String host, int port, long lastHeartbeat, Set<String> services) {
            this.id = id;
            this.host = host;
            this.port = port;
            this.lastHeartbeat = lastHeartbeat;
            this.services = services;
        }

        public String getId() { return id; }
        public String getHost() { return host; }
        public int getPort() { return port; }
        public long getLastHeartbeat() { return lastHeartbeat; }
        public Set<String> getServices() { return services; }
        public void setLastHeartbeat(long lastHeartbeat) { this.lastHeartbeat = lastHeartbeat; }
    }
}